import asyncio
import os

from dotenv import load_dotenv
from langchain.schema.runnable import RunnablePassthrough
from langchain_core.messages import BaseMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_deepseek import ChatDeepSeek

load_dotenv(override=True)
deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")

# 初始化 deepseek
llm = ChatDeepSeek(
    model="deepseek-chat",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=deepseek_api_key,
)
llm_prompt = PromptTemplate(
    template="你是一位资深的Langchain开发专家，回答请严谨。"
)

chain = (
        {"query": RunnablePassthrough()}
        | llm_prompt
        | llm
        | StrOutputParser()
)

questions = [
    "解释一下 LangChain 是什么",
    "LangChain 里的 Chain 指什么",
    "LangChain 的一个版本是什么时候",
    "LangChain 的发起者是谁",
    "LangChain 目前最新的版本是什么"
]
# 并发量控制-如果问题很多，想限制同时发出的请求数（任务都会执行，只是限制了同一时刻的最大并发）
sem = asyncio.Semaphore(3)  # 最多 5 个并发


async def ask_one(q: str) -> BaseMessage:
    """单条异步请求"""
    async with sem:
        return await llm.ainvoke(q)


async def batch_ask():
    """并发请求所有问题，全部完成后统一返回"""
    tasks = [asyncio.create_task(ask_one(q)) for q in questions]
    answers = await asyncio.gather(*tasks)  # 顺序与 questions 一一对应
    return dict(zip(questions, answers))


if __name__ == "__main__":
    print("--------------llm请求开始-前")
    result = asyncio.run(batch_ask())
    for q, a in result.items():
        print(f"【问】{q}\n【答】{a}\n")
    print("--------------llm请求开始-后")
